import numpy as np
from scipy.interpolate import interp1d


""" A module to plot observation HI CDDF data. """


OmegaM_wmap7 = 0.272
OmegaM_owls = 0.238
OmegaM_mill = 0.25



class ObsDat:
    pass



def return_noterdaeme_09( OmegaM_sim = OmegaM_owls ):
    """ set Noterdaeme_fH1_xy_09 
    http://labs.adsabs.harvard.edu/ui/abs/2009A&A...505.1087N """

    fname = 'noterdaeme_fH1_xy_09.txt'
    dat = np.loadtxt( fname )

    notr = ObsDat()
    notr.Nlo = dat[:,0]
    notr.Nhi = dat[:,1]
    notr.f = dat[:,2]

    notr.yerr = dat[:,3]

    notr.Nmi = (notr.Nlo + notr.Nhi)/2
    notr.xerr = (notr.Nhi - notr.Nlo)/2
    
    # adjust for cosmology absorption distance
    # fac = sqrt(OmegaM_Sim / OmegaM_Not) 

    OmegaM_Not = 0.3
    fac = np.sqrt( OmegaM_sim / OmegaM_Not )
    notr.f = np.log10( 10**notr.f * fac )

    return notr



def plot_noterdaeme_09( ax, data, alpha=0.1, xoff=0.0,
                        color='green', marker=None, label='NPLS09' ):

    if not marker == None:
        ax.scatter( data.Nmi+xoff, 
                    data.f, 
                    marker=marker, 
                    alpha=alpha, 
                    label=label, 
                    color=color ) 

    ax.errorbar( data.Nmi+xoff, 
                 data.f, 
                 xerr=data.xerr, 
                 yerr=data.yerr, 
                 ls='None', 
                 alpha=alpha, 
                 color=color,
                 capsize=0,
                 elinewidth=1.5,
                 ecolor=color)





def return_noterdaeme_12( OmegaM_sim = OmegaM_owls ):
    """ set Noterdaeme_fH1_xy_12 
    http://adsabs.harvard.edu/abs/2012A%26A...547L...1N """

    fname = 'noterdaeme_fH1_xy_12.txt'
    dat = np.loadtxt( fname )

    notr = ObsDat()
    print notr
    notr.Nlo = dat[:,0]
    notr.Nhi = dat[:,1]
    notr.f_uncorr = dat[:,2]
    notr.f = dat[:,3]

    notr.yerr = dat[:,4]

    notr.Nmi = (notr.Nlo + notr.Nhi)/2
    notr.xerr = (notr.Nhi - notr.Nlo)/2
    
    # adjust for cosmology absorption distance
    # fac = sqrt(OmegaM_Sim / OmegaM_Not) 

    OmegaM_Not = 0.27
    fac = np.sqrt( OmegaM_sim / OmegaM_Not )
    notr.f = np.log10( 10**notr.f * fac )

    return notr



def plot_noterdaeme_12( ax, data, alpha=0.1, xoff=0.0,
                        color='green', marker=None, label='NPCP12' ):

    if not marker == None:
        ax.scatter( data.Nmi+xoff, 
                    data.f, 
                    marker=marker, 
                    alpha=alpha, 
                    label=label, 
                    color=color ) 

    ax.errorbar( data.Nmi+xoff, 
                 data.f, 
                 xerr=data.xerr, 
                 yerr=data.yerr, 
                 ls='None', 
                 alpha=alpha, 
                 color=color,
                 capsize=0,
                 elinewidth=1.5,
                 ecolor=color)







def return_prochaska_09( OmegaM_sim = OmegaM_owls ):

    """ set Prochaska_fH1_xy_09 
    http://labs.adsabs.harvard.edu/ui/abs/2009ApJ...696.1543P """

    fname = 'prochaska_fH1_xy_09.txt'
    dat = np.loadtxt( fname )

    proc = ObsDat()
    proc.Nlo = dat[:,0]
    proc.Nhi = dat[:,1]
    proc.f = dat[:,2]
    proc.yerr_lo = dat[:,3]
    proc.yerr_hi = dat[:,4]

    proc.Nmi = (proc.Nlo + proc.Nhi)/2
    proc.xerr_lo = proc.Nmi - proc.Nlo
    proc.xerr_hi = proc.Nhi - proc.Nmi

    proc.xerr = np.zeros( (2,proc.Nmi.size) )
    proc.yerr = np.zeros( (2,proc.Nmi.size) )

    proc.xerr[0,:] = proc.xerr_lo
    proc.xerr[1,:] = proc.xerr_hi

    proc.yerr[0,:] = proc.yerr_lo
    proc.yerr[1,:] = proc.yerr_hi

    # adjust for cosmology absorption distance
    # fac = sqrt(OmegaM_Sim / OmegaM_Not) 

    OmegaM_Pro = 0.3
    fac = np.sqrt( OmegaM_sim / OmegaM_Pro )
    proc.f = np.log10( 10**proc.f * fac )

    return proc




def plot_prochaska_09( ax, data, alpha=0.1, xoff=0.0,
                       color='blue', marker=None, label='PW09' ):

    if not marker == None:
        ax.scatter( data.Nmi[0:-3]+xoff, 
                    data.f[0:-3], 
                    marker=marker, 
                    alpha=alpha, 
                    label=label, 
                    color=color ) 

    ax.errorbar( data.Nmi[0:-3]+xoff, 
                 data.f[0:-3], 
                 xerr=data.xerr[:,0:-3], 
                 yerr=data.yerr[:,0:-3], 
                 ls='None', 
                 alpha=alpha, 
                 color=color,
                 ecolor=color,
                 capsize=0,
                 elinewidth=1.5)





def return_omeara_07_esi( OmegaM_sim = OmegaM_owls ):
    """ set O'meara ESI 
    http://labs.adsabs.harvard.edu/ui/abs/2007ApJ...656..666O """


    fname = 'omeara_ESI.txt'
    dat = np.loadtxt( fname )

    omeaESI = ObsDat()
    omeaESI.Nmi = dat[:,0]
    omeaESI.f = dat[:,1]
    omeaESI.xerr_lo = dat[:,2]
    omeaESI.xerr_hi = dat[:,3]
    omeaESI.yerr_lo = dat[:,4]
    omeaESI.yerr_hi = dat[:,5]

    omeaESI.xerr = np.zeros( (2,omeaESI.Nmi.size) )
    omeaESI.yerr = np.zeros( (2,omeaESI.Nmi.size) )
    omeaESI.xerr[0,:] = omeaESI.xerr_lo
    omeaESI.xerr[1,:] = omeaESI.xerr_hi
    omeaESI.yerr[0,:] = omeaESI.yerr_lo
    omeaESI.yerr[1,:] = omeaESI.yerr_hi


    # adjust for cosmology absorption distance
    # fac = sqrt(OmegaM_Sim / OmegaM_Ome) 

    OmegaM_Ome = 0.3
    fac = np.sqrt( OmegaM_sim / OmegaM_Ome )
    omeaESI.f = np.log10( 10**omeaESI.f * fac )

    return omeaESI



def plot_omeara_07_esi( ax, data, alpha=1.0, xoff=0.0,
                        color='green', marker=None, label='OPBP07e' ):
    
    if not marker==None:
        ax.scatter( data.Nmi+xoff, 
                    data.f, 
                    marker=marker, 
                    alpha=alpha, 
                    label=label, 
                    color=color ) 
    
    ax.errorbar( data.Nmi+xoff, 
                 data.f, 
                 xerr=data.xerr, 
                 yerr=data.yerr, 
                 ls='None', 
                 alpha=alpha, 
                 color=color,
                 capsize=0,
                 elinewidth=1.5 )




def return_omeara_07_mike( OmegaM_sim = OmegaM_owls ):
    """ set O'meara MIKE 
    http://labs.adsabs.harvard.edu/ui/abs/2007ApJ...656..666O """

    fname = 'omeara_MIKE.txt'
    dat = np.loadtxt( fname )

    omeaMIKE = ObsDat()
    omeaMIKE.Nmi = dat[:,0]
    omeaMIKE.f = dat[:,1]
    omeaMIKE.xerr_lo = dat[:,2]
    omeaMIKE.xerr_hi = dat[:,3]
    omeaMIKE.yerr_lo = dat[:,4]
    omeaMIKE.yerr_hi = dat[:,5]
    
    omeaMIKE.xerr = np.zeros( (2,omeaMIKE.Nmi.size) )
    omeaMIKE.yerr = np.zeros( (2,omeaMIKE.Nmi.size) )
    omeaMIKE.xerr[0,:] = omeaMIKE.xerr_lo
    omeaMIKE.xerr[1,:] = omeaMIKE.xerr_hi
    omeaMIKE.yerr[0,:] = omeaMIKE.yerr_lo
    omeaMIKE.yerr[1,:] = omeaMIKE.yerr_hi
    
    # adjust for cosmology absorption distance
    # fac = sqrt(OmegaM_Sim / OmegaM_Ome) 

    OmegaM_Ome = 0.3
    fac = np.sqrt( OmegaM_sim / OmegaM_Ome )
    omeaMIKE.f = np.log10( 10**omeaMIKE.f * fac )

    return omeaMIKE



def plot_omeara_07_mike( ax, data, alpha=1.0, xoff=0.0,
                         color='green', marker=None, label='OPBP07m' ):
    
    if not marker==None:
        ax.scatter( data.Nmi+xoff, 
                    data.f, 
                    marker=marker, 
                    alpha=alpha, 
                    label=label, 
                    color=color ) 
    
    ax.errorbar( data.Nmi+xoff, 
                 data.f, 
                 xerr=data.xerr, 
                 yerr=data.yerr, 
                 ls='None', 
                 alpha=alpha, 
                 color=color,
                 capsize=0,
                 elinewidth=1.5 )



def plot_prochaska_09_slopes( ax, beta_4=-6.0, color='grey', 
                              alt_color='orange', alpha=1.0):

    """ Plot data from http://adsabs.harvard.edu/abs/2010ApJ...718..392P """

    OmegaM_Pro = 0.3
    fac = np.sqrt( OmegaM_owls / OmegaM_Pro )
    #fac = np.sqrt( OmegaM_wmap7 / OmegaM_Pro )

    # DLA Branch
    # log NHI > 20.3
    #===========================================================
    N_dla = 10**20.3
    N_d = 10**21.75
    N_cut = 10**22.2

    beta_3 = -1.8
    k_dla = 7.0e-25

    l_dla3 =  k_dla / ( N_d**beta_3 * (beta_3 + 1) ) * \
             ( N_d**(beta_3+1) - N_dla**(beta_3+1) )

    l_dla4 = -k_dla / ( N_d**beta_4 * (beta_4 + 1) ) * \
             N_d**(beta_4+1)

    l_dla = l_dla3 + l_dla4


    xx = np.log10( np.linspace( N_d, N_cut ) )
    yy = np.log10( k_dla * ( 10**xx/N_d )**beta_4 )
    #ax.plot( xx, yy, color=color, alpha=alpha )


    xx = np.log10( np.linspace( N_dla, N_d ) )
    yy = np.log10( k_dla * ( 10**xx/N_d )**beta_3 )
    #ax.plot( xx, yy, color=color, alpha=alpha )


    # Super Lyman Limit System Branch
    # 19.0 < log NHI < 20.3
    #===========================================================

    l_slls = 0.2



    # set up x-axis
    #--------------------------------
    xlims_slls = np.array( [19.0,20.3] )
    logNlo = xlims_slls[0]
    logNhi = xlims_slls[1]
    Nlo = 10**logNlo
    Nhi = 10**logNhi
    xx = logNlo + np.linspace( 0.0, 1.0 ) * (logNhi-logNlo) 

    xx_fill = np.array( [logNlo, logNhi] )

    yy_fill_hi = xx_fill.copy()
    yy_fill_lo = xx_fill.copy()


    # y-axis Beta = -1.0
    #--------------------------------
    beta_slls = -1.2 + 0.2
    k_slls = l_slls  / (np.log(Nhi) - np.log(Nlo))
    ylo = np.log10( k_slls * Nlo**beta_slls )
    yhi = np.log10( k_slls * Nhi**beta_slls )
    yy = ylo + np.linspace( 0.0, 1.0 ) * (yhi - ylo)        
    yy = np.log10( 10**yy * fac )    
    #ax.plot( xx, yy, color=alt_color, alpha=alpha )

    yy_fill_hi[0] = yy[0]
    yy_fill_lo[1] = yy[-1]

    # y-axis Beta = -1.2
    #--------------------------------
    beta_slls = -1.2
    k_slls = l_slls * (beta_slls+1) / (Nhi**(beta_slls+1) - Nlo**(beta_slls+1))
    ylo = np.log10( k_slls * Nlo**beta_slls )
    yhi = np.log10( k_slls * Nhi**beta_slls )
    yy = ylo + np.linspace( 0.0, 1.0 ) * (yhi - ylo)        
    yy = np.log10( 10**yy * fac )    
    #ax.plot( xx, yy, color=color, alpha=alpha )



    # y-axis Beta = -1.4
    #--------------------------------
    beta_slls = -1.2 - 0.2
    k_slls = l_slls * (beta_slls+1) / (Nhi**(beta_slls+1) - Nlo**(beta_slls+1))
    ylo = np.log10( k_slls * Nlo**beta_slls )
    yhi = np.log10( k_slls * Nhi**beta_slls )
    yy = ylo + np.linspace( 0.0, 1.0 ) * (yhi - ylo)        
    yy = np.log10( 10**yy * fac )    
    #ax.plot( xx, yy, color=alt_color, alpha=alpha )

    yy_fill_hi[1] = yy[-1]
    yy_fill_lo[0] = yy[0]


    ax.fill_between( xx_fill, yy_fill_lo, yy_fill_hi,
                      color=color, alpha=0.6 )

    #print 'xx_fill: ', xx_fill
    #print 'yy_fill_lo: ', yy_fill_lo
    #print 'yy_fill_hi: ', yy_fill_hi


    # Lyman Limit System Branch
    # N_plls < log NHI < 19.0
    #
    # and Lyman-Alpha Forest
    # 14.5 < log NHI < N_plls
    #===========================================================
    logNlo = 14.51
    logNhi = 19.0

    Npts = 100
    xx_fill = np.linspace( logNlo, logNhi, Npts )
    yy_fill_hi = xx_fill.copy() - 1000.0
    yy_fill_lo = xx_fill.copy() + 1000.0


    f19s = [ -20.05, -20.05, -20.05,
             -20.25, -20.25, -20.25, -20.25,
             -19.85, -19.85, -19.85, -19.85, 
             -19.85, -19.85, -19.85, -19.85, 
             -19.85, -19.85, -19.85, -19.85,
             -19.85 ]

    beta_lls = [ 0.8, 0.8, 0.8,
                 0.9, 1.3, 1.3, 1.3,
                 0.8, 0.8, 0.8, 0.8,
                 0.8, 0.8, 0.8, 0.8,
                 0.8, 0.8, 0.8, 0.8, 
                 0.8 ]
                 
    Nplls = [ 17.1, 17.3, 17.5,
              17.5, 15.0, 15.2, 15.4,
              15.2, 15.4, 15.5, 15.7,
              15.9, 16.1, 16.2, 16.4,
              16.6, 16.8, 17.0, 17.1,
              17.3 ]

    beta_plls = [ 2.0, 1.9, 1.9,
                  1.9, 3.5, 3.0, 2.6,
                  5.2, 4.3, 3.7, 3.3,
                  3.0, 2.7, 2.5, 2.4,
                  2.2, 2.1, 2.0, 1.9,
                  1.9 ]


    for ii in range( len(f19s) ):

        ylow = f19s[ii]
        xlims = [ Nplls[ii], 19.0 ]
        yhigh = ylow + (xlims[1]-xlims[0]) * beta_lls[ii]
        xx = xlims[0] + np.linspace( 0.0, 1.0 ) * (xlims[1]-xlims[0])
        yy = yhigh - np.linspace( 0.0, 1.0 ) * (yhigh-ylow)
        yy = np.log10( 10**yy * fac )
        #ax.plot( xx, yy, color=alt_color, alpha=0.6 )

        ff = interp1d(xx, yy, 'linear')
        for ix,xx_tmp in enumerate( xx_fill ):
            if xx_tmp > xx[0] and xx_tmp < xx[-1]:
                ff_tmp = ff(xx_tmp)
                if ff_tmp > yy_fill_hi[ix]:
                    yy_fill_hi[ix] = ff_tmp
                if ff_tmp < yy_fill_lo[ix]:
                    yy_fill_lo[ix] = ff_tmp


        xlims = [ 14.5, Nplls[ii] ]
        ylow = yhigh
        yhigh = ylow + (xlims[1]-xlims[0]) * beta_plls[ii]
        xx = xlims[0] + np.linspace( 0.0, 1.0 ) * (xlims[1]-xlims[0])
        yy = yhigh - np.linspace( 0.0, 1.0 ) * (yhigh-ylow)
        yy = np.log10( 10**yy * fac )
        #ax.plot( xx, yy, color=alt_color, alpha=0.6 )

        ff = interp1d(xx, yy, 'linear')
        for ix,xx_tmp in enumerate( xx_fill ):
            if xx_tmp > xx[0] and xx_tmp < xx[-1]:
                ff_tmp = ff(xx_tmp)
                if ff_tmp > yy_fill_hi[ix]:
                    yy_fill_hi[ix] = ff_tmp
                if ff_tmp < yy_fill_lo[ix]:
                    yy_fill_lo[ix] = ff_tmp



    # plot fill
    #-------------------------------------------
    yy_fill_lo[-1] = yy_fill_lo[-2]
    yy_fill_hi[-1] = yy_fill_hi[-2]

    ax.fill_between( xx_fill, yy_fill_lo, yy_fill_hi,
                      color=color, alpha=0.6 )


    # prefered model
    #-------------------------------------------
    ylow = -20.05
    xlims = [17.3,19.0]
    yhigh = ylow + (xlims[1]-xlims[0]) * (0.8)
    xx = xlims[0] + np.linspace( 0.0, 1.0 ) * (xlims[1]-xlims[0])
    yy = yhigh - np.linspace( 0.0, 1.0 ) * (yhigh-ylow)
    yy = np.log10( 10**yy * fac )
    #ax.plot( xx, yy, color=color, alpha=alpha )

    xlims = [14.5,17.3]
    ylow = yhigh
    yhigh = ylow + (xlims[1]-xlims[0]) * (1.9)
    xx = xlims[0] + np.linspace( 0.0, 1.0 ) * (xlims[1]-xlims[0])
    yy = yhigh - np.linspace( 0.0, 1.0 ) * (yhigh-ylow)
    yy = np.log10( 10**yy * fac )
    #ax.plot( xx, yy, color=color, alpha=alpha )






if __name__ == "__main__":

    import matplotlib.pyplot as plt

    notr09 = return_noterdaeme_09()
    notr12 = return_noterdaeme_12()
    proc09 = return_prochaska_09()
    omea07e = return_omeara_07_esi()
    omea07m = return_omeara_07_mike()




    fig = plt.figure(figsize=(8,8))
    ax = fig.add_subplot(111)

    plot_noterdaeme_09(ax,notr09,color='lime')
    plot_noterdaeme_12(ax,notr12,color='green')
    plot_prochaska_09(ax,proc09,color='red')
    plot_omeara_07_esi(ax,omea07e,color='blue')
    plot_omeara_07_mike(ax,omea07m,color='blue')

    plot_prochaska_09_slopes(ax)
